#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-09-18

"""
【节点名称】：
    ImgDirReaderNode
【依赖项安装】：
    pip install spirems
【订阅类型】：
    无
【发布类型】：
    sensor_msgs::CompressedImage （输出图像）
【构造参数说明】：
    parameter_file (str): 全局参数文件
【节点参数】：
    resize (list): 图像重采样分辨率
    fps (int): 输出帧率
【备注】：
    无
"""
import os
import threading
import time
import cv2
import uuid
from typing import Union
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, QoS
from spirecv.base.BaseNode import BaseNode


class ImgDirReaderNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown_emit: bool = True
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=True
        )
        self.sms_shutdown_emit = sms_shutdown_emit
        self.client_id = str(uuid.uuid4()).replace('-', '_')
        self.img_i_queue = Queue()
        self.img_i = 0
        self._next_reader = Subscriber(
            '/' + job_name + '/launch_next', 'std_msgs::Boolean', self.launch_next,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._image_writer = Publisher(
            '/' + job_name + '/sensor/image_raw', 'sensor_msgs::CompressedImage',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self.frame_id = 0
        self.image_dir = self.get_param("image_dir", "G:/dataset/BB_Object365/images/val")
        files_and_directories = os.listdir(self.image_dir)
        self.image_fns = []
        for img_name in files_and_directories:
            if os.path.splitext(img_name)[-1] in ['.jpg', '.jpeg', '.png']:
                self.image_fns.append(os.path.join(self.image_dir, img_name))
        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_writer.kill()
        self._next_reader.kill()

    def launch_next(self, msg: dict = None):
        if (isinstance(msg, dict) and msg['data']) or msg is None:
            if self.img_i < len(self.image_fns):
                self.img_i_queue.put(self.img_i)
                self.img_i += 1

    def run(self):
        while self.is_running():
            img_i = self.img_i_queue.get(block=True)
            if img_i is None:
                break

            img = cv2.imread(self.image_fns[img_i])

            msg = cvimg2sms(img)
            msg['img_id'] = img_i
            msg['img_total'] = len(self.image_fns)
            msg['file_name'] = os.path.basename(self.image_fns[img_i])
            msg['client_id'] = self.client_id
            self._image_writer.publish(msg)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    node = ImgDirReaderNode('detection')
